package rtpps

import (
	"sync"
	"time"
	
	"github.com/pion/rtp"
	"github.com/yapingcat/gomedia/mpeg2"
	"gogs.smarteye.com/smarteye/gortc/pkg/format/rtpmpeg4video"
)

type Decoder struct {
	*rtpmpeg4video.Decoder

	mutex    sync.Mutex
	demuxer  *mpeg2.PSDemuxer
	cidToPid map[mpeg2.PS_STREAM_TYPE]uint8
}

// Init initializes the decoder.
func (d *Decoder) Init() error {
	if err := d.Decoder.Init(); err != nil {
		return err
	}

	d.demuxer = mpeg2.NewPSDemuxer()
	return nil
}

func (d *Decoder) Decode(pkt *rtp.Packet, onFrame func(data []byte, cid mpeg2.PS_STREAM_TYPE, pts time.Duration)) error {
	d.mutex.Lock()
	defer d.mutex.Unlock()

	d.demuxer.OnFrame = func(data []byte, cid mpeg2.PS_STREAM_TYPE, pts uint64, dts uint64) {
		tPTS := time.Duration(pts) * time.Millisecond
		onFrame(data, cid, tPTS)
	}

	data, err := d.Decoder.Decode(pkt)
	if err != nil {
		return err
	}

	if err = d.demuxer.Input(data); err != nil {
		return err
	}

	return nil
}
